﻿module Physics.Chapter5.CollisionResponse

open Casanova
open Casanova.Input
open Casanova.Core
open Casanova.Coroutines
open Casanova.Utilities
open Casanova.Math
open Casanova.Game
open Casanova.Drawing
open Casanova.StandardLibrary
open Casanova.StandardLibrary.Core
open Physics.Chapter5.RigidBody
open Physics.Chapter5.Algebra
open Physics.Chapter5.LookupMap

type FrictionConstraint = {
  Body1               : Ref<RigidBody>
  Body2               : Ref<RigidBody>
  T                   : Vector3<1>
  N                   : Vector3<1>
  B                   : Vector3<1>
  P                   : Vector3<m>
  WeightFraction      : float32<1>
}

type DistanceConstraint = {
  Body1               : Ref<RigidBody>
  Body2               : Ref<RigidBody>
  R1                  : Vector3<m>
  R2                  : Vector3<m>
  L                   : float32<m>
}

type CollisionConstraint = {
  Body1               : Ref<RigidBody>
  Body2               : Ref<RigidBody>
  N                   : Vector3<1>
  P                   : Vector3<m>
  Depth               : float32<m>
}

let mutable private prev_x = [||]

let collision_response (collision_constraints:List<CollisionConstraint>)
                       (friction_constraints:List<FrictionConstraint>)
                       (distance_constraints:List<DistanceConstraint>)
                       (bodies:RigidBody[]) (composite_bodies:List<Ref<RigidBody>[]>) 
                       (dt:float32<s>) =
  if dt < 1.0f<s> / 30.0f then
    let dt_left = 900.0f * (1.0f<s> / 30.0f - dt) |> int
    System.Threading.Thread.Sleep(dt_left)
  let dt = 1.0f<s> / 30.0f

  // tuples of the type i,j (i < j), a_i, b_i, a_j, b_j, k
  // a_i, b_i are the constraint vectors for body i
  // a_j, b_j are the constraint vectors for body j
  // k is the pushing factor, to allow the force to 
  //   perform work and restore the constraint if it 
  //   is broken
  let constraints = 
    [|
      for c in collision_constraints do
        let b_i,b_j = !c.Body1, !c.Body2

        if !b_i.Asleep && b_i.Immovable |> not then
          do b_i.Asleep := false
        if !b_j.Asleep && b_j.Immovable |> not then
          do b_j.Asleep := false

        let x_i,x_j = !b_i.Position, !b_j.Position
        let r_i,r_j = c.P - x_i, c.P - x_j
        let n = c.N * 1.0f<_>
        let k = 
          if abs c.Depth < 0.01f<m> then 0.001f 
          else c.Depth * 0.1f<_>

        let j1,j2 = -n, -Vector3<1>.Cross(r_i, n)
        let j3,j4 =  n,  Vector3<1>.Cross(r_j, n)

        let j1,j2,j3,j4 = 
          if b_i.MassInv = 0.0f<_> then
            Vector3.Zero, Vector3.Zero, j3, j4
          elif b_j.MassInv = 0.0f<_> then
            j1, j2, Vector3.Zero, Vector3.Zero
          else
            j1, j2, j3, j4

        yield b_i, b_j, 
              j1, j2,
              j3, j4,
              k,
              -1000.0f, 0.0f

      for c in friction_constraints do
        let b_i,b_j = !c.Body1, !c.Body2

        let G_i,G_j =
          (if b_i.MassInv <> 0.0f<kg^-1> then
              Vector3<N>.Create(0.0f, float32(-9.81f / b_i.MassInv), 0.0f)
            else 
              Vector3<N>.Zero),
          (if b_j.MassInv <> 0.0f<kg^-1> then
              Vector3<N>.Create(0.0f, float32(-9.81f / b_j.MassInv), 0.0f)
            else 
              Vector3<N>.Zero)

        let F_rel = (!b_j.F + G_i - !b_i.F - G_j) * c.WeightFraction
        let x_i,x_j = !b_i.Position, !b_j.Position
        let r_i,r_j = c.P - x_i, c.P - x_j
        let n = c.N * 1.0f<_>
        let F_n = Vector3<1>.Dot(F_rel, n) |> abs
        let mu = (b_i.FrictionCoefficient + b_j.FrictionCoefficient)
        let u = c.T * 1.0f<_>
        let k = 0.0f
        let j1,j2 = -u, -Vector3<1>.Cross(r_i, u)
        let j3,j4 =  u,  Vector3<1>.Cross(r_j, u)
        if Vector3<1>.Dot(u, Vector3.UnitY) |> abs < 0.1f then
          yield b_i, b_j, 
                j1, j2,
                j3, j4,
                k,
                -mu * F_n, mu * F_n
        let u = c.B * 1.0f<_>
        let k = 0.0f
        let j1,j2 = -u, -Vector3<1>.Cross(r_i, u)
        let j3,j4 =  u,  Vector3<1>.Cross(r_j, u)
        if Vector3<1>.Dot(u, Vector3.UnitY) |> abs < 0.1f then
          yield b_i, b_j, 
                j1, j2,
                j3, j4,
                k,
                -mu * F_n, mu * F_n

      for dc in distance_constraints do
        let b_i,b_j = !dc.Body1, !dc.Body2

        let p_i,p_j = b_i.FromLocalSpace(dc.R1),
                      b_j.FromLocalSpace(dc.R2)
        let x_i,x_j = !b_i.Position, !b_j.Position
        let r_i,r_j = p_i - x_i, p_j - x_j
        let d =  p_j - p_i
        let L = dc.L
        let d_len = d.Length

        if (abs(d_len - L) > 0.1f<m>) || !b_i.Asleep |> not || !b_j.Asleep |> not then
          do b_i.Asleep := false
          do b_j.Asleep := false
          let k = 0.5f * (L - d_len)
          let k = k * 0.1f<_>

          let j1,j2 = -d, -Vector3<1>.Cross(r_i, d)
          let j3,j4 =  d,  Vector3<1>.Cross(r_j, d)

          yield b_i, b_j, 
                j1, j2,
                j3, j4,
                k,
                -10000.0f, 10000.0f
    |]


  // Generate islands
  let constraint_links = [| for b in bodies do yield ResizeArray() |]
  let explored_constraints = [| for c in constraints do yield false |]
  for i = 0 to constraints.Length - 1 do
    let c = constraints.[i]
    let b_i,b_j,_,_,_,_,_,_,_ = c
    let b_i,b_j = b_i.Index,b_j.Index
    constraint_links.[b_i].Add(b_j,i)
    constraint_links.[b_j].Add(b_i,i)

  let mutable min_free = 0
  let mutable island_id = 1
  let island_mapping = [| for b in bodies do yield -11111 |]
  let islands = Stack<ResizeArray<_> * ResizeArray<_> * Var<int> * int>()
  while min_free < bodies.Length do
    let fst (x,y,z,w) = x
    let snd (x,y,z,w) = y
    let trd (x,y,z,w) = z
    islands.Push(ResizeArray(),ResizeArray(), var 0, island_id)
    let frontier = Stack(10)
    frontier.Push min_free
    let explore island_id i = 
      if bodies.[i].Immovable then
        if island_mapping.[i] <> -island_id then
          (islands.Peek() |> fst).Add(bodies.[i])             
          island_mapping.[i] <- -island_id
      else
        if island_mapping.[i] = -11111 then
          island_mapping.[i] <- island_id
          frontier.Push i
          (islands.Peek() |> fst).Add(bodies.[i])
          if bodies.[i].Awake then
            (islands.Peek() |> trd) := !(islands.Peek() |> trd) + 1
        for _,c in constraint_links.[i] do
          if explored_constraints.[c] |> not then
            explored_constraints.[c] <- true
            (islands.Peek() |> snd).Add(constraints.[c])
    do explore island_id min_free
    while frontier.Count > 0 do
      let i = frontier.Pop()
      for j,c in constraint_links.[i] do
        explore island_id j
    // add all static constraints between bodies inside the island
    let island_bodies,island_constraints,island_awake,_ = islands.Peek()
    if !island_awake > 0 then
      for b in island_bodies do
        if b.Immovable then
          for b_j,c_i in constraint_links.[b.Index] do
            let c = constraints.[c_i]
            if island_mapping.[b_j] = island_id && explored_constraints.[c_i] |> not then
              explored_constraints.[c_i] <- true
              island_constraints.Add c
    while min_free < bodies.Length && 
          (island_mapping.[min_free] <> -11111 || bodies.[min_free].Immovable) do
      min_free <- min_free + 1
    island_id <- island_id + 1

  for bodies, constraints, awake, island_id in islands do
    if !awake > 0 then
      for i = 0 to bodies.Length - 1 do
        let b = bodies.[i]
        do b.Index <- i
        do b.IslandId <- island_id
        if (!b.F).Length > 0.0f<N> || (!b.T).Length > 0.0f<N * m> then
          b.Asleep := false
        if !b.Asleep |> not && b.MassInv <> 0.0f<kg^-1> then
          let G = Vector3<N>.Create(0.0f, float32(-9.81f / b.MassInv), 0.0f)
          b.F := !b.F + G

      if constraints.Length = 0 then 
        for b_i in bodies do
          if b_i.Immovable |> not then
            let v2 = !b_i.Velocity + b_i.MassInv * !b_i.F * dt
            let w2 = !b_i.AngularVelocity + b_i.InverseI * !b_i.T * dt * 1.0f<_>
            b_i.Velocity := v2
            b_i.AngularVelocity := w2
            if (!b_i.Velocity).Length <= 0.001f<m/s> && (!b_i.AngularVelocity).Length <= 0.001f<m^2/s> then
              b_i.Asleep := true
            b_i.Position := !b_i.Position + !b_i.Velocity * dt
            b_i.Rotation := !b_i.Rotation + b_i.RotationVelocity * (float32 dt)
            b_i.Rotation := Microsoft.Xna.Framework.Quaternion.Normalize(!b_i.Rotation)
      else
        let b = 
          [|
            for b_i,b_j,dv_i,dw_i,dv_j,dw_j,k,x_min,x_max in constraints do
              let e = float32 k
              let v_i,w_i,v_j,w_j = !b_i.Velocity, !b_i.AngularVelocity, !b_j.Velocity, !b_j.AngularVelocity
              let jv = Vector3<1>.Dot(v_i, dv_i) + Vector3<1>.Dot(w_i, dw_i * 1.0f<_>) +
                        Vector3<1>.Dot(v_j, dv_j) + Vector3<1>.Dot(w_j, dw_j * 1.0f<_>)

              let f_i,t_i,f_j,t_j = !b_i.F, !b_i.T, !b_j.F, !b_j.T
              let a_i,al_i,a_j,al_j = f_i * b_i.MassInv, b_i.InverseI * t_i,
                                      f_j * b_j.MassInv, b_j.InverseI * t_j
              let jmf = Vector3<1>.Dot(dv_i, a_i) + Vector3<1>.Dot(dw_i, al_i * 1.0f<_>) +
                          Vector3<1>.Dot(dv_j, a_j) + Vector3<1>.Dot(dw_j, al_j * 1.0f<_>)
              yield e / dt * 1.0f<_> - jv * 1.0f<_> - 1.0f<_> * jmf * dt
          |] |> Vector.Create

        let JMInverse = 
          let n = bodies.Length
          [|
            for b_i,b_j,dv_i,dw_i,dv_j,dw_j,k,x_min,x_max in constraints do
              let v_i,w_i,v_j,w_j = !b_i.Velocity, !b_i.AngularVelocity, !b_j.Velocity, !b_j.AngularVelocity

              let f_i,t_i,f_j,t_j = !b_i.F, !b_i.T, !b_j.F, !b_j.T
              let a_i,al_i,a_j,al_j = f_i * b_i.MassInv, b_i.InverseI * t_i,
                                      f_j * b_j.MassInv, b_j.InverseI * t_j
              let jmf = Vector3<1>.Dot(dv_i, a_i) + Vector3<1>.Dot(dw_i, al_i * 1.0f<_>) +
                          Vector3<1>.Dot(dv_j, a_j) + Vector3<1>.Dot(dw_j, al_j * 1.0f<_>)

              yield b_i.Index,b_j.Index,dv_i * b_i.MassInv,b_i.InverseI * dw_i,dv_j * b_j.MassInv,b_j.InverseI * dw_j
          |]

        let body_constraints = 
          [|
            for b in bodies do
              yield LookupMap.Create constraints.Length
          |]

        let constraints_intersections =
          [|
            for c in constraints do
              yield LookupMap.Create constraints.Length
          |]

        for c_i = 0 to constraints.Length - 1 do
          let b_i,b_j,_,_,_,_,_,_,_ = constraints.[c_i]
          if b_i.Immovable |> not then
            body_constraints.[b_i.Index].Add c_i
            for c_i' in body_constraints.[b_i.Index].Entries do
              constraints_intersections.[c_i].Add c_i'
              constraints_intersections.[c_i'].Add c_i
          if bodies.[b_j.Index].Immovable |> not then
            body_constraints.[b_j.Index].Add c_i
            for c_i' in body_constraints.[b_j.Index].Entries do
              constraints_intersections.[c_i].Add c_i'
              constraints_intersections.[c_i'].Add c_i

        let A = 
          let m = constraints.Length    
          [|
            for x = 0 to m - 1 do
              yield 
                [|
//                  for y = 0 to m - 1 do
                  for y = 0 to constraints_intersections.[x].Entries.Count - 1 do
                    let y = constraints_intersections.[x].Entries.[y]

                    // sum row x of JMInverse and row y of J
                    let i0_x,i1_x,a_i0_x,b_i0_x,a_i1_x,b_i1_x = JMInverse.[x]
                    let i0_y,i1_y,a_i0_y,b_i0_y,a_i1_y,b_i1_y,_,_,_ = constraints.[y]
                    let i0_y,i1_y = i0_y.Index,i1_y.Index

                    let sum = var 0.0f
                    let inline check_and_sum (i_x,a_i_x,b_i_x) (i_y,a_i_y,b_i_y) =
                        if i_x = i_y then
                          sum := !sum + Vector3<1>.Dot(a_i_x, a_i_y) +
                                        Vector3<1>.Dot(b_i_x, b_i_y)

                    let (i_x,a_i_x,b_i_x) = (i0_x,a_i0_x,b_i0_x)
                    let (i_y,a_i_y,b_i_y) = (i0_y,a_i0_y,b_i0_y)
                    do check_and_sum (i_x,a_i_x,b_i_x) (i_y,a_i_y * 1.0f<_>,b_i_y * 1.0f<_>)
                    let (i_y,a_i_y,b_i_y) = (i1_y,a_i1_y,b_i1_y)
                    do check_and_sum (i_x,a_i_x,b_i_x) (i_y,a_i_y * 1.0f<_>,b_i_y * 1.0f<_>)

                    let (i_x,a_i_x,b_i_x) = (i1_x,a_i1_x,b_i1_x)
                    let (i_y,a_i_y,b_i_y) = (i0_y,a_i0_y,b_i0_y)
                    do check_and_sum (i_x,a_i_x,b_i_x) (i_y,a_i_y * 1.0f<_>,b_i_y * 1.0f<_>)
                    let (i_y,a_i_y,b_i_y) = (i1_y,a_i1_y,b_i1_y)
                    do check_and_sum (i_x,a_i_x,b_i_x) (i_y,a_i_y * 1.0f<_>,b_i_y * 1.0f<_>)
                  
                    if !sum <> 0.0f then
                      yield y,!sum
                |]
          |]
        let A = SparseMatrix.Create(A, constraints.Length, constraints.Length)

        for bc in body_constraints do 
          LookupMap.Delete(bc)
        for cc in constraints_intersections do 
          LookupMap.Delete(cc)

        let x0 = 
          if prev_x.Length <> constraints.Length then
            [| for c in constraints do yield 0.0f |] |> Algebra.Vector.Create
          else
            prev_x |> Algebra.Vector.Create
        let x_min = [| for i,j,a_i,b_i,a_j,b_j,k,x_min,x_max in constraints do yield x_min |] |> Algebra.Vector.Create
        let x_max = [| for i,j,a_i,b_i,a_j,b_j,k,x_min,x_max in constraints do yield x_max |] |> Algebra.Vector.Create
        let lambda = projected_gauss_seidel A b x0 x_min x_max
        do prev_x <- lambda.V

        // let Fc = J.Transpose * lambda
        let Fc = 
          [|
            for b in bodies do
              yield Vector3.Zero
          |]
        let Tc = 
          [|
            for b in bodies do
              yield Vector3.Zero
          |]

        for (i,j,a_i,b_i,a_j,b_j,_,_,_),lambda in Seq.zip constraints lambda.V do
          let i,j = i.Index,j.Index
          let F_i,T_i = a_i * lambda, b_i * lambda
          let F_j,T_j = a_j * lambda, b_j * lambda
          Fc.[i] <- Fc.[i] + F_i
          Fc.[j] <- Fc.[j] + F_j
          Tc.[i] <- Tc.[i] + T_i
          Tc.[j] <- Tc.[j] + T_j
      
        let V2 = 
          [|
            for b in bodies do
              let Fc,Tc = Fc.[b.Index],Tc.[b.Index]
              let F = !b.F * dt + Fc * 1.0f<_>
              let T = !b.T * dt + Tc * 1.0f<_>
              let v2 = !b.Velocity + b.MassInv * F
              let w2 = !b.AngularVelocity + b.InverseI * T * 1.0f<_>

              yield! v2 |> vec_to_seq
              yield! w2 |> vec_to_seq
          |] |> Vector.Create

        for i = 0 to bodies.Length - 1 do
          let b_i = bodies.[i]
          let v = Vector3(V2.[i * 3 * 2 + 0],
                          V2.[i * 3 * 2 + 1],
                          V2.[i * 3 * 2 + 2])
          let w = Vector3(V2.[i * 3 * 2 + 3],
                          V2.[i * 3 * 2 + 4],
                          V2.[i * 3 * 2 + 5])
          b_i.Velocity := v * 1.0f<_>
          b_i.AngularVelocity := w * 1.0f<_>
          if (!b_i.Velocity).Length <= 0.1f<m/s> && (!b_i.AngularVelocity).Length <= 0.1f<m^2/s> then
            b_i.Asleep := true
          else
            if b_i.MassInv <> 0.0f<_> then
              b_i.Position := !b_i.Position + !b_i.Velocity * dt
              b_i.Rotation := !b_i.Rotation + b_i.RotationVelocity * (float32 dt)
              b_i.Rotation := Microsoft.Xna.Framework.Quaternion.Normalize(!b_i.Rotation)

  for c in composite_bodies do
    let avg_velocity = (c |> Seq.sumBy (fun b -> !(!b).Velocity)) / (float32 c.Length)
    let avg_angular_velocity = (c |> Seq.sumBy (fun b -> !(!b).AngularVelocity)) / (float32 c.Length)
    
    if avg_velocity.Length <= (float32 c.Length) * 0.1f<m/s> && avg_angular_velocity.Length <= (float32 c.Length) * 0.1f<m^2/s> then
      for b in c do
        (!b).Asleep := true
        